
#include "co.h"
#include <stdlib.h>
#include <setjmp.h>
#include <stdint.h>
#include <string.h>

#define STACK_SIZE (64 * 1024)

static inline void
stack_switch_call(void *sp, void *entry, uintptr_t arg) {
    asm volatile (
#if __x86_64__
        "movq %0, %%rsp; movq %2, %%rdi; jmp *%1"
          :
          : "b"((uintptr_t)sp),
            "d"(entry),
            "a"(arg)
          : "memory"
#else
            "movl %0, %%esp; movl %2, 4(%0); jmp *%1"
            :
            : "b"((uintptr_t)sp - 8),
    "d"(entry),
    "a"(arg)
            : "memory"
#endif
            );
}


enum co_status {
    CO_NEW = 1, // new, not have been used
    CO_RUNNING = 2, // have been used
    CO_WAITING = 3, // wait
    CO_DEAD = 4, // have ended, but not have been freed
};



struct co {
    const char* name;
    void (*func)(void *);   //func start entrance
    void *arg;            // arguments

    enum co_status status;
    struct co* waiter;   // the co wait

    jmp_buf context;
    uint8_t stack[STACK_SIZE];

    void *stackptr;

    struct co* next;
};

struct co* current;
struct co* co_main;



__attribute__((constructor)) void co_init() {
    co_main = (struct co*)malloc(sizeof(struct co));
    co_main->name = "main";
    co_main->stackptr = co_main->stack + sizeof(co_main->stack);
    co_main->status = CO_RUNNING;
    memset(co_main->stack, 0, sizeof(co_main->stack));
    co_main->next = co_main;

    current = co_main;

}

struct co *co_start(const char *name, void (*func)(void *), void *arg) {// create a co
    struct co*  res = (struct co*) malloc (sizeof(struct co));
    res->name = name;
    res->func = func;
    res->arg = arg;
    res->status = CO_NEW;
    res->stackptr = res->stack+sizeof(res->stack);

    memset(res->stack,0, sizeof(res->stack));


    if(current->next == current){
        res->next = co_main;
        co_main->next = res;
        return res;
    }


    // add it to the cycle_cos
    struct co* part = current;
    while(part->next != co_main){
        part = part->next;
    }


    part->next = res;
    res->next = co_main;

    return res;

}

void wrapper(){
    current->status=CO_RUNNING;
    current->func(current->arg);

    //a long time ...
    current->status=CO_DEAD;

    if(current->waiter!=NULL){
        current->waiter->status=CO_RUNNING;
        current->waiter=NULL;
    }
    co_yield();
}

void co_yield() {

    //  assert(current);
    int val = setjmp(current->context);
    if(val == 0)  // by co_yield
    {
        struct co* coNext = current->next;
        while(coNext->status == CO_WAITING || coNext->status == CO_DEAD){
            coNext = coNext->next;
        }
        current = coNext;
        if(coNext->status == CO_RUNNING){
            longjmp(current->context, 1);
        }
        else{
            if(sizeof(void*)==4)//  如果指针大小为4字节（32位系统），使用函数调用切换栈
            {
                stack_switch_call(coNext->stackptr,wrapper,(uintptr_t)NULL);
            }
            else{ // 如果指针大小不为4字节（64位系统），使用内联汇编直接切换栈
                asm volatile("mov %0,%%rsp"::"b"((uintptr_t)coNext->stackptr));
                wrapper();
            }

        }
    }
    else{ // val == 1, by long jmp
        return;
    }
}

void co_wait(struct co *co) {

    if(co->status != CO_DEAD){
        current->status = CO_WAITING;
        co->waiter = current;
        co_yield();
        current->status = CO_RUNNING;
    }



    //delete
    struct co* part = current;

    while (part->next != co)
    {
        part = part->next;
    }

    part->next = part->next->next;
    free(co);

}